import os
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.embeddings import QianfanEmbeddingsEndpoint
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.llms import QianfanLLMEndpoint

os.environ["QIANFAN_AK"] = "7ZITkqmufR2g6MSLQUqsYAXB"
os.environ["QIANFAN_SK"] = "n2YLOoakGri5LfsK8Fxsax86tdeLEfpm"

WEB_URL = "https://zhuanlan.zhihu.com/p/85289282" # 使用WebBaseLoader加载HTML
loader = WebBaseLoader(WEB_URL)
docs = loader.load()
# 导入千帆向量模型
embeddings = QianfanEmbeddingsEndpoint()
# 导入递归字符文本分割器
text_splitter = RecursiveCharacterTextSplitter(chunk_size = 384, chunk_overlap = 0, separators=["\n\n", "\n", " ", "", "。", "，"])
# 导入文本
documents = text_splitter.split_documents(docs)
# 存入向量数据库
vector = Chroma.from_documents(documents, embeddings)
# 创建提示词模板
prompt = ChatPromptTemplate.from_template("""使用下面的语料来回答本模板最末尾的问题。如果你不知道问题的答案，直接回答 "我不知道"，禁止随意编造答案。         
为了保证答案尽可能简洁，你的回答必须不超过三句话，你的回答中不可以带有星号。         
请注意！在每次回答结束之后，你都必须接上 "感谢你的提问" 作为结束语         
以下是一对问题和答案的样例：             
请问：秦始皇的原名是什么             
秦始皇原名嬴政。感谢你的提问。                  
以下是语料： <context> {context} </context>  Question: {input}""")
#创建千帆LLM模型
llm = QianfanLLMEndpoint()
#我们设置一个链，该链接受一个问题和检索到的文档并生成一个答案。
document_chain = create_stuff_documents_chain(llm, prompt)
#使用检索器动态选择最相关的文档并将其传递。
retriever = vector.as_retriever()
retrieval_chain = create_retrieval_chain(retriever, document_chain)
#调用这个链了。这将返回一个字典 - 来自 LLM 的响应在键中answer
response = retrieval_chain.invoke({
    "input": "朱元璋时期中国多少人口?"
})
print(response["answer"])